"""
整合database、prompt、llm_chat
"""
from ai.llm.prompt_factory import PromptTemplateFactory
from ai.llm.database_helper import milvusDatabase
from ai.llm.llm_chat_client import qianfanLLM
from ai.llm.param_builder import MilvusParam


def testQuery(query):
    # 1、问题向量化
    embeddings = qianfanLLM.embed_query(query)
    # 2、检索向量数据库
    db_param = MilvusParam.getUserParam()
    ans = milvusDatabase.searchEmbedding([embeddings], db_param)
    # 3、预选答案准备
    examples = []
    for item in ans[0]:
        examples.append({"example": item.fields["desc"]})
    # 4、构建Prompt
    prompt = PromptTemplateFactory.testPromptTemplate(query, examples)
    # 5、调用LLM
    llm_result = qianfanLLM.invoke(prompt)
    return {"prompt": prompt.format_prompt().text, "result": llm_result.content}


def boatNormQuery(query):
    # 1、问题向量化
    embeddings = qianfanLLM.embed_query(query)
    # 2、检索向量数据库
    db_param = MilvusParam.getBoatNormParam()
    ans = milvusDatabase.searchEmbedding([embeddings], db_param)
    # 3、预选答案准备
    examples = []
    for item in ans[0]:
        examples.append({"example": item.fields["content"]})
    # 4、构建Prompt
    prompt = PromptTemplateFactory.boatNormPromptTemplate(query, examples)
    # 5、调用LLM
    llm_result = qianfanLLM.invoke(prompt)
    return {"prompt": prompt.format_prompt().text, "result": llm_result.content}


if __name__ == '__main__':
    # 检查 NPU 是否可用
    # try:
    #     from transformers.utils import is_torch_npu_available
    #     if is_torch_npu_available():
    #         print("NPU is available.")
    #     else:
    #         print("NPU is not available.")
    # except ImportError:
    #     print("Please install Transformers library with NPU support.")
    result = boatNormQuery("滚装船检验的标准是什么？")
    print(result)
